{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DDJwQPZcupab"
   },
   "source": [
    "# EECS 498-007/598-005 Assignment 6: Generative Adversarial Networks \n",
    "\n",
    "Before we start, please put your name and UMID in following format\n",
    "\n",
    ": Firstname LASTNAME, #00000000   //   e.g.) Justin JOHNSON, #12345678"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2KMxqLt1h2kx"
   },
   "source": [
    "**Your Answer:**   \n",
    "Hello WORLD, #XXXXXXXX"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gtaO2qGf_mdG",
    "tags": [
     "pdf-title"
    ]
   },
   "source": [
    "# Generative Adversarial Networks (GANs)\n",
    "\n",
    "So far in this course, all the applications of neural networks that we have explored have been **discriminative models** that take an input and are trained to produce a labeled output. This has ranged from straightforward classification of image categories to sentence generation (which was still phrased as a classification problem, our labels were in vocabulary space and we’d learned a recurrence to capture multi-word labels). In this notebook, we will expand our repetoire, and build **generative models** using neural networks. Specifically, we will learn how to build models which generate novel images that resemble a set of training images.\n",
    "\n",
    "### What is a GAN?\n",
    "\n",
    "In 2014, [Goodfellow et al.](https://arxiv.org/abs/1406.2661) presented a method for training generative models called Generative Adversarial Networks (GANs for short). In a GAN, we build two different neural networks. Our first network is a traditional classification network, called the **discriminator**. We will train the discriminator to take images, and classify them as being real (belonging to the training set) or fake (not present in the training set). Our other network, called the **generator**, will take random noise as input and transform it using a neural network to produce images. The goal of the generator is to fool the discriminator into thinking the images it produced are real.\n",
    "\n",
    "We can think of this back and forth process of the generator ($G$) trying to fool the discriminator ($D$), and the discriminator trying to correctly classify real vs. fake as a minimax game:\n",
    "$$\\underset{G}{\\text{minimize}}\\; \\underset{D}{\\text{maximize}}\\; \\mathbb{E}_{x \\sim p_\\text{data}}\\left[\\log D(x)\\right] + \\mathbb{E}_{z \\sim p(z)}\\left[\\log \\left(1-D(G(z))\\right)\\right]$$\n",
    "where $z \\sim p(z)$ are the random noise samples, $G(z)$ are the generated images using the neural network generator $G$, and $D$ is the output of the discriminator, specifying the probability of an input being real. In [Goodfellow et al.](https://arxiv.org/abs/1406.2661), they analyze this minimax game and show how it relates to minimizing the Jensen-Shannon divergence between the training data distribution and the generated samples from $G$.\n",
    "\n",
    "To optimize this minimax game, we will aternate between taking gradient *descent* steps on the objective for $G$, and gradient *ascent* steps on the objective for $D$:\n",
    "1. update the **generator** ($G$) to minimize the probability of the __discriminator making the correct choice__. \n",
    "2. update the **discriminator** ($D$) to maximize the probability of the __discriminator making the correct choice__.\n",
    "\n",
    "While these updates are useful for analysis, they do not perform well in practice. Instead, we will use a different objective when we update the generator: maximize the probability of the **discriminator making the incorrect choice**. This small change helps to allevaiate problems with the generator gradient vanishing when the discriminator is confident. This is the standard update used in most GAN papers, and was used in the original paper from [Goodfellow et al.](https://arxiv.org/abs/1406.2661). \n",
    "\n",
    "In this assignment, we will alternate the following updates:\n",
    "1. Update the generator ($G$) to maximize the probability of the discriminator making the incorrect choice on generated data:\n",
    "$$\\underset{G}{\\text{maximize}}\\;  \\mathbb{E}_{z \\sim p(z)}\\left[\\log D(G(z))\\right]$$\n",
    "2. Update the discriminator ($D$), to maximize the probability of the discriminator making the correct choice on real and generated data:\n",
    "$$\\underset{D}{\\text{maximize}}\\; \\mathbb{E}_{x \\sim p_\\text{data}}\\left[\\log D(x)\\right] + \\mathbb{E}_{z \\sim p(z)}\\left[\\log \\left(1-D(G(z))\\right)\\right]$$\n",
    "\n",
    "### What else is there?\n",
    "Since 2014, GANs have exploded into a huge research area, with massive [workshops](https://sites.google.com/site/nips2016adversarial/), and [hundreds of new papers](https://github.com/hindupuravinash/the-gan-zoo). Compared to other approaches for generative models, they often produce the highest quality samples but are some of the most difficult and finicky models to train (see [this github repo](https://github.com/soumith/ganhacks) that contains a set of 17 hacks that are useful for getting models working). Improving the stabiilty and robustness of GAN training is an open research question, with new papers coming out every day! For a more recent tutorial on GANs, see [here](https://arxiv.org/abs/1701.00160). There is also some even more recent exciting work that changes the objective function to Wasserstein distance and yields much more stable results across model architectures: [WGAN](https://arxiv.org/abs/1701.07875), [WGAN-GP](https://arxiv.org/abs/1704.00028).\n",
    "\n",
    "\n",
    "GANs are not the only way to train a generative model! For other approaches to generative modeling check out the [deep generative model chapter](http://www.deeplearningbook.org/contents/generative_models.html) of the Deep Learning [book](http://www.deeplearningbook.org). Another popular way of training neural networks as generative models is Variational Autoencoders (co-discovered [here](https://arxiv.org/abs/1312.6114) and [here](https://arxiv.org/abs/1401.4082)). Variatonal autoencoders combine neural networks with variationl inference to train deep generative models. These models tend to be far more stable and easier to train but currently don't produce samples that are as pretty as GANs.\n",
    "\n",
    "Here's an example of what your outputs from the 3 different models you're going to train should look like... note that GANs are sometimes finicky, so your outputs might not look exactly like this... this is just meant to be a *rough* guideline of the kind of quality you can expect:\n",
    "\n",
    "![caption](https://web.eecs.umich.edu/~justincj/teaching/eecs498/assets/a6/gan_outputs_pytorch.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UzKziWkw_mdK",
    "tags": [
     "pdf-ignore"
    ]
   },
   "source": [
    "## Setup code\n",
    "Before getting started, we need to run some boilerplate code to set up our environment, same as previous assignments. You'll need to rerun this setup code each time you start the notebook.\n",
    "\n",
    "First, run this cell load the autoreload extension. This allows us to edit .py source files, and re-import them into the notebook for a seamless editing and debugging experience."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GOondkpep9v5"
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sPrpYW0Fp_4N"
   },
   "source": [
    "### Google Colab Setup\n",
    "Next we need to run a few commands to set up our environment on Google Colab. If you are running this notebook on a local machine you can skip this section.\n",
    "\n",
    "Run the following cell to mount your Google Drive. Follow the link, sign in to your Google account (the same account you used to store this notebook!) and copy the authorization code into the text box that appears below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "hNe2zxXIqx2Z"
   },
   "outputs": [],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/content/drive')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uaoUvjWXqHo9"
   },
   "source": [
    "Now recall the path in your Google Drive where you uploaded this notebook, fill it in below. If everything is working correctly then running the folowing cell should print the filenames from the assignment:\n",
    "\n",
    "```\n",
    "['eecs598', 'gan.py', 'generative_adversarial_networks.ipynb', 'a6_helper.py', 'vae.py', 'variational_autoencoders.ipynb']\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-gEkt4gsqICp"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# TODO: Fill in the Google Drive path where you uploaded the assignment\n",
    "# Example: If you create a 2020FA folder and put all the files under A5 folder, then '2020FA/A5'\n",
    "GOOGLE_DRIVE_PATH_AFTER_MYDRIVE = None\n",
    "GOOGLE_DRIVE_PATH = os.path.join('drive', 'My Drive', GOOGLE_DRIVE_PATH_AFTER_MYDRIVE)\n",
    "print(os.listdir(GOOGLE_DRIVE_PATH))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "at01aCfTqVpZ"
   },
   "source": [
    "Once you have successfully mounted your Google Drive and located the path to this assignment, run th following cell to allow us to import from the `.py` files of this assignment. If it works correctly, it should print the message:\n",
    "\n",
    "```\n",
    "Hello from gan.py!\n",
    "Hello from a6_helper.py!\n",
    "```\n",
    "\n",
    "as well as the last edit time for the file `gan.py`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "h10YGmfmqWfH"
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(GOOGLE_DRIVE_PATH)\n",
    "\n",
    "import time, os\n",
    "os.environ[\"TZ\"] = \"US/Eastern\"\n",
    "time.tzset()\n",
    "\n",
    "from gan import hello_gan\n",
    "hello_gan()\n",
    "\n",
    "from a6_helper import hello_helper\n",
    "hello_helper()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "78rH-FoC_mdL",
    "tags": [
     "pdf-ignore"
    ]
   },
   "outputs": [],
   "source": [
    "import math\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import init\n",
    "import torchvision\n",
    "import torchvision.transforms as T\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data import sampler\n",
    "import torchvision.datasets as dset\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.gridspec as gridspec\n",
    "\n",
    "from eecs598.grad import rel_error\n",
    "from eecs598.utils import reset_seed\n",
    "\n",
    "%matplotlib inline\n",
    "plt.rcParams['figure.figsize'] = (10.0, 8.0) # set default size of plots\n",
    "plt.rcParams['image.interpolation'] = 'nearest'\n",
    "plt.rcParams['image.cmap'] = 'gray'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FudOhR-d_mdg",
    "tags": [
     "pdf-ignore"
    ]
   },
   "source": [
    "Note that if CUDA is not enabled, `torch.cuda.is_available()` will return False and this notebook will fallback to CPU mode.\n",
    "\n",
    "The global variables `dtype` and `device` will control the data types throughout this assignment.\n",
    "\n",
    "We will be using `torch.float = torch.float32` for all operations.\n",
    "\n",
    "Please refer to https://pytorch.org/docs/stable/tensor_attributes.html#torch-dtype for more details about data types."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "blz1sXlkIL_q",
    "tags": [
     "pdf-ignore-input"
    ]
   },
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "  print('Good to go!')\n",
    "else:\n",
    "\n",
    "  print('Please set GPU via Edit -> Notebook Settings.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Fa68wtAMQJKr"
   },
   "source": [
    " Now we'll import some helper functions from the `a6_helper.py` we've provided to help us visualize the MNIST dataset, verify the neural networks you implement, and initialize weights in the `torch.nn` modules you'll write."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CzZnoR2pZ5rY"
   },
   "outputs": [],
   "source": [
    "from a6_helper import show_images, count_params, initialize_weights\n",
    "\n",
    "dtype = torch.float\n",
    "device = 'cuda'\n",
    "answers = {}\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "K6MV195M_mdR",
    "tags": [
     "pdf-ignore"
    ]
   },
   "source": [
    "## Dataset\n",
    " GANs are notoriously finicky with hyperparameters, and also require many training epochs. In order to make this assignment approachable, we will be working on the MNIST dataset, which is 60,000 training and 10,000 test images. Each picture contains a centered image of white digit on black background (0 through 9). This was one of the first datasets used to train convolutional neural networks and it is fairly easy -- a standard CNN model can easily exceed 99% accuracy. \n",
    "\n",
    "To simplify our code here, we will use the PyTorch MNIST wrapper, which downloads and loads the MNIST dataset. See the [documentation](https://github.com/pytorch/vision/blob/master/torchvision/datasets/mnist.py) for more information about the interface. The default parameters will take 5,000 of the training examples and place them into a validation dataset. The data will be saved into a folder called `MNIST_data`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "hVIGFRTz_mdS",
    "tags": [
     "pdf-ignore"
    ]
   },
   "outputs": [],
   "source": [
    "batch_size = 128\n",
    "NOISE_DIM = 96\n",
    "\n",
    "print('download MNIST if not exist')\n",
    "\n",
    "mnist_train = dset.MNIST('./MNIST_data', train=True, download=True,\n",
    "                           transform=T.ToTensor())\n",
    "loader_train = DataLoader(mnist_train, batch_size=batch_size,\n",
    "                          shuffle=True, drop_last=True, num_workers=2)\n",
    "\n",
    "\n",
    "imgs = loader_train.__iter__().next()[0].view(batch_size, 784)\n",
    "show_images(imgs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fZvir2Cm_mdV"
   },
   "source": [
    "## Random Noise\n",
    "The first step is to generate uniform noise from -1 to 1 with shape `[batch_size, noise_dim]`\n",
    "\n",
    "Hint: use `torch.rand`.\n",
    "\n",
    "Implement `sample_noise` and verify all tests pass below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NdUieMiy_mdZ"
   },
   "outputs": [],
   "source": [
    "from gan import sample_noise\n",
    "reset_seed(0)\n",
    "\n",
    "\n",
    "batch_size = 3\n",
    "noise_dim = 4\n",
    "\n",
    "z = sample_noise(batch_size, noise_dim)\n",
    "assert z.shape == (batch_size, noise_dim)\n",
    "assert torch.is_tensor(z)\n",
    "assert torch.all(z >= -1.0) and torch.all(z <= 1.0)\n",
    "assert torch.any(z < 0.0) and torch.any(z > 0.0)\n",
    "print('All tests passed!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "16F3IOHK_mdj"
   },
   "source": [
    "# Discriminator\n",
    "Our first step is to build a discriminator. All fully connected layers should include bias terms. The architecture is:\n",
    " * Fully connected layer with input size 784 and output size 256\n",
    " * LeakyReLU with alpha 0.01\n",
    " * Fully connected layer with input_size 256 and output size 256\n",
    " * LeakyReLU with alpha 0.01\n",
    " * Fully connected layer with input size 256 and output size 1\n",
    " \n",
    "Recall that the Leaky ReLU nonlinearity computes $f(x) = \\max(\\alpha x, x)$ for some fixed constant $\\alpha$; for the LeakyReLU nonlinearities in the architecture above we set $\\alpha=0.01$.\n",
    " \n",
    "The output of the discriminator should have shape `[batch_size, 1]`, and contain real numbers corresponding to the scores that each of the `batch_size` inputs is a real image. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6O2hnRu9_mdo"
   },
   "source": [
    "Implement `discriminator` in `gan.py` and test your solution by running the cell below. Be sure to use `nn.Sequential` for this model definition and all future models in this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "D0bfhJhp_mdp"
   },
   "outputs": [],
   "source": [
    "from gan import discriminator\n",
    "\n",
    "def test_discriminator(true_count=267009):\n",
    "  model = discriminator()\n",
    "  cur_count = count_params(model)\n",
    "  print(cur_count)\n",
    "  if cur_count != true_count:\n",
    "    print('Incorrect number of parameters in discriminator. Check your achitecture.')\n",
    "  else:\n",
    "    print('Correct number of parameters in discriminator.')     \n",
    "\n",
    "test_discriminator()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gIorZ7EK_mds"
   },
   "source": [
    "# Generator\n",
    "Now to build the generator network:\n",
    " * Fully connected layer from noise_dim to 1024\n",
    " * `ReLU`\n",
    " * Fully connected layer with size 1024 \n",
    " * `ReLU`\n",
    " * Fully connected layer with size 784\n",
    " * `TanH` (to clip the image to be in the range of [-1,1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CCOQsxPR_mdx"
   },
   "source": [
    "Implement `generator` in `gan.py` and test your solution by running the cell below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-3Og3IjU_mdz"
   },
   "outputs": [],
   "source": [
    "from gan import generator\n",
    "\n",
    "def test_generator(true_count=1858320):\n",
    "  model = generator(4)\n",
    "  cur_count = count_params(model)\n",
    "  print(cur_count)\n",
    "  if cur_count != true_count:\n",
    "    print('Incorrect number of parameters in generator. Check your achitecture.')\n",
    "  else:\n",
    "    print('Correct number of parameters in generator.')\n",
    "\n",
    "test_generator()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ayjC9wo8_md1"
   },
   "source": [
    "# GAN Loss\n",
    "\n",
    "Compute the generator and discriminator loss. The generator loss is:\n",
    "$$\\ell_G  =  -\\mathbb{E}_{z \\sim p(z)}\\left[\\log D(G(z))\\right]$$\n",
    "and the discriminator loss is:\n",
    "$$ \\ell_D = -\\mathbb{E}_{x \\sim p_\\text{data}}\\left[\\log D(x)\\right] - \\mathbb{E}_{z \\sim p(z)}\\left[\\log \\left(1-D(G(z))\\right)\\right]$$\n",
    "Note that these are negated from the equations presented earlier as we will be *minimizing* these losses.\n",
    "\n",
    "For the purpose of these equations, we assume that the output from the discriminator is a real number in the range $0 < D(x) < 1$ which results from squashing the raw score from the discriminator through a sigmoid function. However for a cleaner and more numerically stable implementation, we have not included the sigmoid in the discriminator architecture above -- instead we will implement the sigmoid as part of the loss function.\n",
    "\n",
    "**HINTS**: You can use the function [`torch.nn.functional.binary_cross_entropy_with_logits`](https://pytorch.org/docs/stable/nn.functional.html#binary-cross-entropy-with-logits) to compute these losses in a numerically stable manner.\n",
    "\n",
    "Given a score $s\\in\\mathbb{R}$ and a label $y\\in\\{0, 1\\}$, the binary cross entropy loss (with logits) is defined as:\n",
    "\n",
    "$$ bce(s, y) = -y * \\log(\\sigma(s)) - (1 - y) * \\log(1 - \\sigma(s)) $$\n",
    "\n",
    "where $\\sigma(s)=1/(1+\\exp(-s))$ is the sigmoid function.\n",
    "\n",
    "A naive implementation of this formula can be numerically unstable, so you should prefer to use the built-in PyTorch implementation.\n",
    "\n",
    "You will also need to compute labels corresponding to real or fake and use the logit arguments to determine their size. Make sure you cast these labels to the correct data type using the global `dtype` variable, for example:\n",
    "\n",
    "`true_labels = torch.ones(size, device=device)`\n",
    "\n",
    "Instead of computing the expectation of $\\log D(G(z))$, $\\log D(x)$ and $\\log \\left(1-D(G(z))\\right)$, we will be averaging over elements of the minibatch, so make sure to combine the loss by averaging instead of summing."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RNjvlGIP_md8"
   },
   "source": [
    "Implement `discriminator_loss` and `generator_loss` in `gan.py` and test your solution by running the cell below.. You should see errors < `1e-7`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vcQsiznb_md9"
   },
   "outputs": [],
   "source": [
    "from gan import discriminator_loss\n",
    "\n",
    "answers['logits_fake'] = torch.tensor(\n",
    "  [-1.80865868,  0.09030055, -0.4428902 , -0.07879368, -0.37655044,\n",
    "    0.32084742, -0.28590837,  1.01376281,  0.99241439,  0.39394346],\n",
    "  dtype=dtype, device=device)\n",
    "answers['d_loss_true'] = torch.tensor(1.8423983904443109, dtype=dtype, device=device)\n",
    "answers['logits_real'] = torch.tensor(\n",
    "  [ 0.93487311, -1.01698916, -0.57304769, -0.88162704, -1.40129389,\n",
    "   -1.45395693, -1.54239755, -0.57273325,  0.98584429,  0.13312152],\n",
    "  dtype=dtype, device=device)\n",
    "\n",
    "def test_discriminator_loss(logits_real, logits_fake, d_loss_true):\n",
    "  d_loss = discriminator_loss(logits_real, logits_fake)\n",
    "  print(\"Maximum error in d_loss: %g\"%rel_error(d_loss_true, d_loss))\n",
    "test_discriminator_loss(answers['logits_real'], answers['logits_fake'],\n",
    "                        answers['d_loss_true'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "c_q2hBzD_meA"
   },
   "outputs": [],
   "source": [
    "from gan import generator_loss\n",
    "\n",
    "answers['g_loss_true'] = torch.tensor(0.771286196423346, dtype=dtype, device=device)\n",
    "\n",
    "def test_generator_loss(logits_fake, g_loss_true):\n",
    "  g_loss = generator_loss(logits_fake)\n",
    "  print(\"Maximum error in g_loss: %g\"%rel_error(g_loss_true, g_loss))\n",
    "\n",
    "test_generator_loss(answers['logits_fake'], answers['g_loss_true'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zE1L2nqy_meD"
   },
   "source": [
    "# Optimizing our loss\n",
    "\n",
    "\n",
    "Next, you'll define a function that returns an `optim.Adam` optimizer for the given model with a 1e-3 learning rate, beta1=0.5, beta2=0.999. We'll use this to construct optimizers for the generators and discriminators for the rest of the notebook in `get_optimizer`.\n",
    "\n",
    "Implement `get_optimizer` in `gan.py` before moving forward in the notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "674H6PyJ_meH"
   },
   "source": [
    "# Training a GAN!\n",
    "\n",
    "We provide you the main training loop... you won't need to change this function, but we encourage you to read through and understand it. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8NVHWpSd_meI",
    "tags": [
     "pdf-ignore"
    ]
   },
   "outputs": [],
   "source": [
    "def run_a_gan(D, G, D_solver, G_solver, discriminator_loss, generator_loss, save_filename, show_every=250, \n",
    "              batch_size=128, noise_size=96, num_epochs=10):\n",
    "  \"\"\"\n",
    "  Train a GAN!\n",
    "  \n",
    "  Inputs:\n",
    "  - D, G: PyTorch models for the discriminator and generator\n",
    "  - D_solver, G_solver: torch.optim Optimizers to use for training the\n",
    "    discriminator and generator.\n",
    "  - discriminator_loss, generator_loss: Functions to use for computing the generator and\n",
    "    discriminator loss, respectively.\n",
    "  - show_every: Show samples after every show_every iterations.\n",
    "  - batch_size: Batch size to use for training.\n",
    "  - noise_size: Dimension of the noise to use as input to the generator.\n",
    "  - num_epochs: Number of epochs over the training dataset to use for training.\n",
    "  \"\"\"\n",
    "  iter_count = 0\n",
    "  for epoch in range(num_epochs):\n",
    "    for x, _ in loader_train:\n",
    "      if len(x) != batch_size:\n",
    "        continue\n",
    "      D_solver.zero_grad()\n",
    "      real_data = x.view(-1, 784).to(device)\n",
    "      logits_real = D(2* (real_data - 0.5))\n",
    "\n",
    "      g_fake_seed = sample_noise(batch_size, noise_size, dtype=real_data.dtype, device=real_data.device)\n",
    "      fake_images = G(g_fake_seed).detach()\n",
    "      logits_fake = D(fake_images)\n",
    "\n",
    "      d_total_error = discriminator_loss(logits_real, logits_fake)\n",
    "      d_total_error.backward()        \n",
    "      D_solver.step()\n",
    "\n",
    "      G_solver.zero_grad()\n",
    "      g_fake_seed = sample_noise(batch_size, noise_size, dtype=real_data.dtype, device=real_data.device)\n",
    "      fake_images = G(g_fake_seed)\n",
    "\n",
    "      gen_logits_fake = D(fake_images)\n",
    "      g_error = generator_loss(gen_logits_fake)\n",
    "      g_error.backward()\n",
    "      G_solver.step()\n",
    "\n",
    "      if (iter_count % show_every == 0):\n",
    "        print('Iter: {}, D: {:.4}, G:{:.4}'.format(iter_count,d_total_error.item(),g_error.item()))\n",
    "        imgs_numpy = fake_images.data.cpu()#.numpy()\n",
    "        show_images(imgs_numpy[0:16])\n",
    "        plt.show()\n",
    "        print()\n",
    "      iter_count += 1\n",
    "    if epoch == num_epochs - 1:\n",
    "      show_images(imgs_numpy[0:16])\n",
    "      plt.savefig(os.path.join(GOOGLE_DRIVE_PATH,save_filename))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Mtn7BrtHOOPB"
   },
   "source": [
    "Now run the cell below to train your first GAN! Your last epoch results will be stored in `fc_gan_results.jpg` for you to submit to the autograder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-Bh2wehE_meR"
   },
   "outputs": [],
   "source": [
    "from gan import get_optimizer\n",
    "reset_seed(0)\n",
    "\n",
    "# Make the discriminator\n",
    "D = discriminator().to(device)\n",
    "\n",
    "# Make the generator\n",
    "G = generator().to(device)\n",
    "\n",
    "# Use the function you wrote earlier to get optimizers for the Discriminator and the Generator\n",
    "D_solver = get_optimizer(D)\n",
    "G_solver = get_optimizer(G)\n",
    "# Run it!\n",
    "run_a_gan(D, G, D_solver, G_solver, discriminator_loss, generator_loss, 'fc_gan_results.jpg')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qOAzglMK_meW",
    "tags": [
     "pdf-ignore"
    ]
   },
   "source": [
    "Well that wasn't so hard, was it? In the iterations in the low 100s you should see black backgrounds, fuzzy shapes as you approach iteration 1000, and decent shapes, about half of which will be sharp and clearly recognizable as we pass 3000."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ujBJ8SLm_meX"
   },
   "source": [
    "# Least Squares GAN\n",
    "We'll now look at [Least Squares GAN](https://arxiv.org/abs/1611.04076), a newer, more stable alernative to the original GAN loss function. For this part, all we have to do is change the loss function and retrain the model. We'll implement equation (9) in the paper, with the generator loss:\n",
    "$$\\ell_G  =  \\frac{1}{2}\\mathbb{E}_{z \\sim p(z)}\\left[\\left(D(G(z))-1\\right)^2\\right]$$\n",
    "\n",
    "and the discriminator loss:\n",
    "\n",
    "$$ \\ell_D = \\frac{1}{2}\\mathbb{E}_{x \\sim p_\\text{data}}\\left[\\left(D(x)-1\\right)^2\\right] + \\frac{1}{2}\\mathbb{E}_{z \\sim p(z)}\\left[ \\left(D(G(z))\\right)^2\\right]$$\n",
    "\n",
    "In these equations, we assume that the output from the discriminator is an unbounded real number $-\\infty < D(x) < \\infty$.\n",
    "\n",
    "\n",
    "**HINTS**: Instead of computing the expectation, we will be averaging over elements of the minibatch, so make sure to combine the loss by averaging instead of summing. When plugging in for $D(x)$ and $D(G(z))$ use the direct output from the discriminator (`scores_real` and `scores_fake`)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "i1fn2OqI_mea"
   },
   "source": [
    "Verify your implementation below, you should see relative errors of < `1e-7` or less"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WGrWDusD_meb"
   },
   "outputs": [],
   "source": [
    "from gan import ls_discriminator_loss\n",
    "from gan import ls_generator_loss\n",
    "\n",
    "answers['d_loss_lsgan_true'] = torch.tensor(1.8770293614440594, dtype=dtype, device=device)\n",
    "answers['g_loss_lsgan_true'] = torch.tensor(0.816954786997558, dtype=dtype, device=device)\n",
    "def test_lsgan_loss(score_real, score_fake, d_loss_true, g_loss_true):\n",
    "  d_loss = ls_discriminator_loss(score_real, score_fake)\n",
    "  g_loss = ls_generator_loss(score_fake)\n",
    "  print(\"Maximum error in d_loss: %g\"%rel_error(d_loss_true, d_loss))\n",
    "  print(\"Maximum error in g_loss: %g\"%rel_error(g_loss_true, g_loss))\n",
    "\n",
    "test_lsgan_loss(answers['logits_real'], answers['logits_fake'],\n",
    "                answers['d_loss_lsgan_true'], answers['g_loss_lsgan_true'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wu8O7N75_mee"
   },
   "source": [
    "Run the following cell to train your model! Your last epoch results will be stored in `ls_gan_results.jpg` for you to submit to the autograder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9aMeH5iR_mef"
   },
   "outputs": [],
   "source": [
    "reset_seed(0)\n",
    "\n",
    "D_LS = discriminator().to(device)\n",
    "G_LS = generator().to(device)\n",
    "\n",
    "D_LS_solver = get_optimizer(D_LS)\n",
    "G_LS_solver = get_optimizer(G_LS)\n",
    "\n",
    "run_a_gan(D_LS, G_LS, D_LS_solver, G_LS_solver, ls_discriminator_loss, ls_generator_loss, 'ls_gan_results.jpg')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ULFkXw00_mej"
   },
   "source": [
    "# Deeply Convolutional GANs\n",
    "In the first part of the notebook, we implemented an almost direct copy of the original GAN network from Ian Goodfellow. However, this network architecture allows no real spatial reasoning. It is unable to reason about things like \"sharp edges\" in general because it lacks any convolutional layers. Thus, in this section, we will implement some of the ideas from [DCGAN](https://arxiv.org/abs/1511.06434), where we use convolutional networks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7pz0MKv434KQ"
   },
   "source": [
    "#### Discriminator\n",
    "We will use a discriminator inspired by the TensorFlow MNIST classification tutorial, which is able to get above 99% accuracy on the MNIST dataset fairly quickly. \n",
    "* Reshape into image tensor (Use `nn.Unflatten`!)\n",
    "* Conv2D: 32 Filters, 5x5, Stride 1\n",
    "* Leaky ReLU(alpha=0.01)\n",
    "* Max Pool 2x2, Stride 2\n",
    "* Conv2D: 64 Filters, 5x5, Stride 1\n",
    "* Leaky ReLU(alpha=0.01)\n",
    "* Max Pool 2x2, Stride 2\n",
    "* Flatten\n",
    "* Fully Connected with output size 4 x 4 x 64\n",
    "* Leaky ReLU(alpha=0.01)\n",
    "* Fully Connected with output size 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KB_tDRgT_mek"
   },
   "outputs": [],
   "source": [
    "from gan import build_dc_classifier\n",
    "\n",
    "data = next(enumerate(loader_train))[-1][0].to(dtype=dtype, device=device)\n",
    "batch_size = data.size(0)\n",
    "b = build_dc_classifier().to(device)\n",
    "data = data.view(-1, 784)\n",
    "out = b(data)\n",
    "print(out.size())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2sNaW7rw_men"
   },
   "source": [
    "Check the number of parameters in your classifier as a sanity check:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Hnzw9vXq_meo"
   },
   "outputs": [],
   "source": [
    "def test_dc_classifer(true_count=1102721):\n",
    "  model = build_dc_classifier()\n",
    "  cur_count = count_params(model)\n",
    "  print(cur_count)\n",
    "  if cur_count != true_count:\n",
    "    print('Incorrect number of parameters in generator. Check your achitecture.')\n",
    "  else:\n",
    "    print('Correct number of parameters in generator.')\n",
    "\n",
    "test_dc_classifer()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RfH2fcBv_mew"
   },
   "source": [
    "#### Generator\n",
    "For the generator, we will copy the architecture exactly from the [InfoGAN paper](https://arxiv.org/pdf/1606.03657.pdf). See Appendix C.1 MNIST. See the documentation for [nn.ConvTranspose2d](https://pytorch.org/docs/stable/nn.html#convtranspose2d). We are always \"training\" in GAN mode. \n",
    "* Fully connected with output size 1024\n",
    "* `ReLU`\n",
    "* BatchNorm\n",
    "* Fully connected with output size 7 x 7 x 128 \n",
    "* `ReLU`\n",
    "* BatchNorm\n",
    "* Reshape into Image Tensor of shape 7 x 7 x 128\n",
    "* Conv2D^T (Transpose): 64 filters of 4x4, stride 2, 'same' padding (use `padding=1`)\n",
    "* `ReLU`\n",
    "* BatchNorm\n",
    "* Conv2D^T (Transpose): 1 filter of 4x4, stride 2, 'same' padding (use `padding=1`)\n",
    "* `TanH`\n",
    "* Should have a 28 x 28 x 1 image, reshape back into 784 vector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qJRCSqvm_mew"
   },
   "outputs": [],
   "source": [
    "from gan import build_dc_generator\n",
    "\n",
    "test_g_gan = build_dc_generator().to(device)\n",
    "test_g_gan.apply(initialize_weights)\n",
    "\n",
    "fake_seed = torch.randn(batch_size, NOISE_DIM, dtype=dtype, device=device)\n",
    "fake_images = test_g_gan.forward(fake_seed)\n",
    "fake_images.size()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uIz2DO7M_me0"
   },
   "source": [
    "Check the number of parameters in your generator as a sanity check:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "a9uu1yy2_me1"
   },
   "outputs": [],
   "source": [
    "def test_dc_generator(true_count=6580801):\n",
    "  model = build_dc_generator(4)\n",
    "  cur_count = count_params(model)\n",
    "  print(cur_count)\n",
    "  if cur_count != true_count:\n",
    "    print('Incorrect number of parameters in generator. Check your achitecture.')\n",
    "  else:\n",
    "    print('Correct number of parameters in generator.')\n",
    "\n",
    "test_dc_generator()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KZ2275lrSZAS"
   },
   "source": [
    "Now, let's train our DC-GAN! Your last epoch results will be stored in `dc_gan_results.jpg` for you to submit to the autograder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FZUM0ygr_me6"
   },
   "outputs": [],
   "source": [
    "reset_seed(0)\n",
    "\n",
    "D_DC = build_dc_classifier().to(device) \n",
    "D_DC.apply(initialize_weights)\n",
    "G_DC = build_dc_generator().to(device)\n",
    "G_DC.apply(initialize_weights)\n",
    "\n",
    "D_DC_solver = get_optimizer(D_DC)\n",
    "G_DC_solver = get_optimizer(G_DC)\n",
    "\n",
    "run_a_gan(D_DC, G_DC, D_DC_solver, G_DC_solver, discriminator_loss, generator_loss, 'dc_gan_results.jpg', num_epochs=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ccGkOUt71NpJ"
   },
   "source": [
    "#### Latent Space Interpolation\n",
    "As we did with our VAE, our final test of our trained GAN model is to perform interpolation in the latent space. We generate random latent vectors $z_0$ and $z_1$, and linearly interplate between them; we run each interpolated vector through the trained generator to produce an image.\n",
    "\n",
    "Each row of the figure below interpolates between two random vectors. For the most part the model should exhibit smooth transitions along each row, demonstrating that the model has learned something nontrivial about the underlying spatial structure of the digits it is modeling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_2ZPwZIDwL0l"
   },
   "outputs": [],
   "source": [
    "S = 12\n",
    "z0 = sample_noise(S, NOISE_DIM, device=device)\n",
    "z1 = sample_noise(S, NOISE_DIM, device=device)\n",
    "w = torch.linspace(0, 1, S, device=device).view(S, 1, 1)\n",
    "z = (w * z0 + (1 - w) * z1).transpose(0, 1).reshape(S * S, NOISE_DIM)\n",
    "\n",
    "x = G_DC(z)\n",
    "show_images(x.data.cpu())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8QydfzS7Y1Dx"
   },
   "source": [
    "## Final Check\n",
    "\n",
    "Make sure all your training results (loss + images) are saved in the notebook. You can run \"Runtime -> Restart and run all...\" to double check before submitting"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qwS9yIh9RfvG"
   },
   "source": [
    "# Submit Your Work\n",
    "\n",
    "After completing all the notebooks for this assignment, run the following cell to create a .zip file for you to download and turn in. Please MANUALLY SAVE every *.ipynb and *.py files before executing the following cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "V_1_-tM5RrWG"
   },
   "outputs": [],
   "source": [
    "from eecs598.submit import make_a6_submission\n",
    "\n",
    "# TODO: Replace these with your actual uniquename and umid\n",
    "uniquename = None\n",
    "umid = None\n",
    "make_a6_submission(GOOGLE_DRIVE_PATH, uniquename, umid)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "generative_adversarial_networks.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
